package com.example.qqviewspringboot.config;

import com.example.qqviewspringboot.websocket.MessageWebSocketHandler;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    private final MessageWebSocketHandler messageWebSocketHandler;

    @Autowired
    public WebSocketConfig(MessageWebSocketHandler messageWebSocketHandler) {
        this.messageWebSocketHandler = messageWebSocketHandler;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(messageWebSocketHandler, "/ws")
                .addInterceptors(new UserIdHandshakeInterceptor())  // 使用自定义的拦截器
                .setAllowedOrigins("*");  // 支持跨域，根据需要调整
    }

    private static class UserIdHandshakeInterceptor implements HandshakeInterceptor {
        @Override
        public boolean beforeHandshake(org.springframework.http.server.ServerHttpRequest request,
                                       org.springframework.http.server.ServerHttpResponse response,
                                       org.springframework.web.socket.WebSocketHandler wsHandler,
                                       java.util.Map<String, Object> attributes) throws Exception {

            // 解析 URL 中的 userId 参数
            String query = request.getURI().getQuery();
            Long userId = null;
            if (query != null && query.contains("userId=")) {
                try {
                    String userIdStr = query.split("=")[1];
                    userId = Long.parseLong(userIdStr);  // 转换为 Long 类型
                } catch (NumberFormatException e) {
                    // 处理异常，如日志记录
                    System.err.println("Invalid userId format");
                }
            }

            if (userId != null) {
                attributes.put("userId", userId);  // 将 userId 存入 attributes
            }
            return true;
        }

        @Override
        public void afterHandshake(org.springframework.http.server.ServerHttpRequest request,
                                   org.springframework.http.server.ServerHttpResponse response,
                                   org.springframework.web.socket.WebSocketHandler wsHandler, Exception exception) {
            // 不需要做任何操作
        }
    }
}
